package types

import (
	"encoding/json"
	"os"
	"sort"
	"strconv"
	"strings"

	llq "github.com/emirpasic/gods/queues/linkedlistqueue"
	"github.com/emirpasic/gods/sets/hashset"
	"github.com/google/uuid"

	"mig/pkg/digester"
	"mig/pkg/utils"
)

type Digest = []byte

// DigestNode represents a node in the digest tree.
type DigestNode struct {
	ID       string
	Digest   []byte
	Parent   *DigestNode
	Children []*DigestNode
}

func NewDummyDigestNode() *DigestNode {
	return &DigestNode{
		ID:     uuid.New().String(),
		Digest: nil,
	}
}

func (d *DigestNode) IsLeaf() bool {
	return len(d.Children) == 0
}

func (d *DigestNode) IsRoot() bool {
	return d.Parent == nil
}

func (d *DigestNode) GetChunkNo() int {
	idx := strings.LastIndex(d.ID, "_")
	chunkNo, _ := strconv.ParseInt(d.ID[idx+1:], 10, 64)
	return int(chunkNo)
}

func (d *DigestNode) Equal(another *DigestNode) bool {
	return utils.EqualSlice(d.Digest, another.Digest)
}

// DigestTree represents a digest tree.
type DigestTree struct {
	Root     *DigestNode
	digester digester.Digester
	size     int
}

func (dt *DigestTree) getLeaves() []*DigestNode {
	if dt.Root == nil {
		return nil
	}
	var leaves []*DigestNode
	q := llq.New()
	q.Enqueue(dt.Root)
	for !q.Empty() {
		v, _ := q.Dequeue()
		root := v.(*DigestNode)
		for _, child := range root.Children {
			q.Enqueue(child)
		}
		if len(root.Children) == 0 {
			leaves = append(leaves, root.Parent)
		}
	}
	return leaves
}

func (dt *DigestTree) Padding(expectedSize int) *DigestTree {
	if dt.size < expectedSize {
		leaves := dt.getLeaves()
		for i := 0; i < expectedSize-dt.size; i++ {
			leaves = append(leaves, NewDummyDigestNode())
		}
		return NewDigestTree(leaves, dt.digester)
	}
	return dt
}

// NewDigestTree creates a new digest tree from the given nodes and digester.
func NewDigestTree(nodes []*DigestNode, digester digester.Digester) *DigestTree {
	root := generateDigestTree(nodes, digester)
	return &DigestTree{
		Root:     root,
		digester: digester,
		size:     len(nodes),
	}
}

func MergeDigestTree(trees []*DigestTree, digester digester.Digester) *DigestTree {
	roots := make([]*DigestNode, 0, len(trees))
	for _, tree := range trees {
		if tree == nil || tree.Root == nil {
			continue
		}
		roots = append(roots, tree.Root)
	}
	return NewDigestTree(roots, digester)
}

// MergeDigestNodes merges multiple digest nodes into one parent node.
func MergeDigestNodes(nodes []*DigestNode, digester digester.Digester) *DigestNode {
	digester.Reset()
	p := &DigestNode{ID: uuid.NewString()}
	for _, node := range nodes {
		digester.Update(node.Digest)
		p.Children = append(p.Children, node)
		node.Parent = p
	}
	p.Digest = digester.Finalize()
	return p
}

// generateDigestTree generates a digest tree from the given nodes and digester.
func generateDigestTree(nodes []*DigestNode, digester digester.Digester) *DigestNode {
	if len(nodes) == 0 {
		return nil
	}
	q := llq.New()
	for _, node := range nodes {
		q.Enqueue(node)
	}
	n := len(nodes)
	for n > 1 {
		for i := 0; i < n; i += 2 {
			v1, _ := q.Dequeue()
			v2, _ := q.Dequeue()
			p := MergeDigestNodes([]*DigestNode{v1.(*DigestNode), v2.(*DigestNode)}, digester)
			q.Enqueue(p)
		}

		if n%2 == 1 {
			v, _ := q.Dequeue()
			vo := v.(*DigestNode)
			vp := &DigestNode{
				ID:       uuid.NewString(),
				Digest:   vo.Digest,
				Children: []*DigestNode{vo},
			}
			vo.Parent = vp
			q.Enqueue(vp)
		}
		n = q.Size()
	}
	root, _ := q.Dequeue()
	return root.(*DigestNode)
}

// NodeJSON defines the JSON format for a node.
type NodeJSON struct {
	ID     string `json:"id"`
	Digest []byte `json:"digest"`
}

// EdgeJSON defines the JSON format for an edge.
type EdgeJSON struct {
	Source string `json:"source"`
	Target string `json:"target"`
}

// DigestTreeJSON defines the JSON format for a digest tree.
type DigestTreeJSON struct {
	Nodes []NodeJSON `json:"nodes"`
	Edges []EdgeJSON `json:"edges"`
}

func FromJSON(data []byte) (*DigestTree, error) {
	jsonData := DigestTreeJSON{}
	if err := json.Unmarshal(data, &jsonData); err != nil {
		return nil, err
	}

	outDegree := map[string]int{}
	parentMap := map[string]string{}
	for _, edge := range jsonData.Edges {
		outDegree[edge.Source]++
		parentMap[edge.Target] = edge.Source
	}

	nodeMap := make(map[string]*DigestNode)
	var leaves []string
	for _, node := range jsonData.Nodes {
		nodeMap[node.ID] = &DigestNode{
			ID:     node.ID,
			Digest: node.Digest,
		}
		if outDegree[node.ID] == 0 {
			leaves = append(leaves, node.ID)
		}
	}

	sort.Strings(leaves)
	q := llq.New()
	for _, id := range leaves {
		q.Enqueue(id)
	}

	inQ := hashset.New()
	var root *DigestNode
	for !q.Empty() {
		v, _ := q.Dequeue()
		id := v.(string)
		parent, exists := parentMap[id]
		if !exists {
			root = nodeMap[id]
			break
		}
		if !inQ.Contains(parent) {
			q.Enqueue(parent)
			inQ.Add(parent)
		}
		parentNode := nodeMap[parent]
		nodeMap[id].Parent = parentNode
		parentNode.Children = append(parentNode.Children, nodeMap[id])
	}

	return &DigestTree{
		Root:     root,
		digester: digester.NewMD5Digester(),
		size:     len(leaves),
	}, nil
}

func (dt *DigestTree) WriteJSON(path string) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()
	_, err = file.Write(([]byte)(dt.ToJSON()))
	return err
}

// ToJSON converts the DigestTree to a JSON string using BFS.
func (dt *DigestTree) ToJSON() string {
	var nodes []NodeJSON
	var edges []EdgeJSON

	queue := llq.New()
	if dt.Root != nil {
		queue.Enqueue(dt.Root)
	}

	for !queue.Empty() {
		node, _ := queue.Dequeue()
		digestNode := node.(*DigestNode)
		nodes = append(nodes, NodeJSON{
			ID:     digestNode.ID,
			Digest: digestNode.Digest,
		})

		for _, child := range digestNode.Children {
			edges = append(edges, EdgeJSON{
				Source: digestNode.ID,
				Target: child.ID,
			})
			queue.Enqueue(child)
		}
	}

	dtJSON := DigestTreeJSON{
		Nodes: nodes,
		Edges: edges,
	}
	bytes, _ := json.Marshal(dtJSON)
	return string(bytes)
}

func (dt *DigestTree) Empty() bool {
	return dt == nil || dt.size == 0
}
